package com.gitee.wsl.ext.coroutines

import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.selects.select

val <T> T.asDeferred: Deferred<T>
    get() = CompletableDeferred(this)


@OptIn(ExperimentalCoroutinesApi::class)
fun <T> Deferred<T>.getCompletedOrNull(): T? {
    return try {
        getCompleted()
    } catch (_: Throwable) {
        null
    }
}

suspend fun <T> awaitAny(vararg deferreds: Deferred<T>): T {
    require(deferreds.isNotEmpty())
    return select { deferreds.forEach { it.onAwait { it } } }
}

suspend fun <T> Collection<Deferred<T>>.awaitAny(): T {
    require(isNotEmpty())
    return select { forEach { it.onAwait { it } } }
}


suspend fun <T> Collection<Deferred<T>>.awaitAnyAndCancelOthers(): T {
    require(isNotEmpty())
    val firstAwaited = select { forEachIndexed { index, deferred -> deferred.onAwait { IndexedValue(index, it) } } }
    val firstAwaitedIndex = firstAwaited.index
    forEachIndexed { index, deferred -> if (index != firstAwaitedIndex) deferred.cancel() }
    return firstAwaited.value
}



/**
 * Awaits for the completion of this [Deferred] without throwing [CancellationException] when completed.
 *
 * If the deferred is completed normally, this function returns the result.
 *
 * If the deferred is cancelled, this function would rethrow a [CancellationException].
 *
 * @param T the type of the deferred value
 * @return the result value of the deferred
 */
@OptIn(ExperimentalCoroutinesApi::class)
suspend inline fun <T> Deferred<T>.smartAwait(): T {
    return try {
        await()
    } catch (e: CancellationException) {
        try {
            getCompleted()
        } catch (_: IllegalStateException) {
            // Rethrow the original cancellation exception
            throw e
        }
    }
}

/**
 * @author ChenFei(chenfei0928@gmail.com)
 * @date 2020-10-23 18:08
 */
@OptIn(ExperimentalCoroutinesApi::class)
suspend inline fun <T> Deferred<T>.blockCompleteOrLaunchAwait(
    scope: CoroutineScope, crossinline block: suspend (T) -> Unit
) {
    if (isCompleted) {
        block(getCompleted())
    } else {
        scope.launch {
            block(await())
        }
    }
}
